#include "Timer.h"

/* maximum number of messages logged within learning period */
#define MAX_MESSAGES 1000

/* max size of neighbour list */
#define MAX_NEIGHBOURS 10
#define MAX_BLACK_FLAGS 10

#define MAX_RATE_INC 20

/* timer constants */
#define GATHER_TIME 30000
#define DETECT_TIME 10000

/* intrusion detection/prevention constants */
#define ALFA_SIZE  1.1
#define BETA_SIZE  0.7
#define ALFA_RATE  2

/* hook states */
#define NO_STATE 0
#define GATHER 1
#define DETECT 2

/* packet headers */
#define HEADER  0x08

generic module HookReceiverP (am_id_t amId){
	provides{
		interface Receive;
		interface Packet;
		interface AMPacket;
	}
	uses{
		interface Receive as ReceiveM;
		interface Packet as PacketM;
		interface AMPacket as AMPacketM;
		interface Timer<TMilli> as GDTimer;
	}
}
implementation{
	int32_t messageNr = 0;
	int16_t sizeNeighbours = 0;
	int16_t sizeTmpList = 0;

	/* determines weather network info is beeing gathered for a new rate calculation */
	bool rateTime = FALSE;
 
	/* messages keept for the learning process */
	message_t messageVec[MAX_MESSAGES];
        
	/* keeps the size of every message for the learning process */
	int16_t sizeVec[MAX_MESSAGES];
	
	/* parameters used by the IPS */
        int16_t avg_rate,avg_size,tmp_rate = 0;

        int16_t state = NO_STATE;

	/* lists of motes with which made contact */
        int16_t neighbours[MAX_NEIGHBOURS];
        int16_t blackFlag[MAX_BLACK_FLAGS];
	int16_t tmpList[MAX_NEIGHBOURS];

	/*prints the message list*/
	 void printMessages(){
                uint16_t i;
                for (i=0;i<messageNr;i++)
                {
                        dbg("HookReceiverP","Message %d \n",i);
                        /* print message data length for micaz messages*/
                        dbg("HookReceiverP","Data length %d \n",sizeVec[i]);
                        /* print source&destination adress for micaz messages*/
                        dbg("HookReceiverP","Source %d \n",call AMPacket.source(&messageVec[i]));
                        dbg("HookReceiverP","Destination %d \n\n",call AMPacket.destination(&messageVec[i]));

                }
        }
	/* Timer actions */
	event void GDTimer.fired(){

                uint16_t i,j;
                am_addr_t address;
		switch(state){

		/* aplication is in learning stage */
                case GATHER :
                        for(i=0;i<messageNr;i++)
                        {
                                avg_size += sizeVec[i];
                                address = call AMPacket.source(&messageVec[i]);
				if(address/16>MAX_NEIGHBOURS){
					 dbg("HookReceiverP","Error neighbour table is not large enough to add node id %d to it\n",address);
					continue;
					}
				else{	
                                	neighbours[address/16] = neighbours[address/16]|(1<<(address%16));
	                        	if(sizeNeighbours<(address/16+1))
						sizeNeighbours = address/16+1;
				}
			}
			
			/* change aplication stage */
			state = DETECT;
                        avg_rate = messageNr;
			avg_size = avg_size /avg_rate;

			/*******print GATHER info********/
/*			printMessages();
			dbg("HookReceiverP","Gather_Neighbours:\n ");

			for(i =0 ;i<sizeNeighbours;i++)
				for (j=0 ;j<16;j++)
					if(((neighbours[i]>>j)&1 )== 1)
						dbg("HookReceiverP","%d\n",(i*16 + j));

			 dbg("HookReceiverP","Average size: %d \n",avg_size);
			 dbg("HookReceiverP","Average rate: %d \n",avg_rate);
*/
			for(i=0;i<MAX_BLACK_FLAGS;i++)
				blackFlag[i] = 0;
                        messageNr =0;
                        call GDTimer.startOneShot(DETECT_TIME);
                break;
		/* aplication is in detection stage */
                case DETECT:
                       
			 messageNr = 0;
			
			if(rateTime){
                                 rateTime = FALSE;
                                 if(tmp_rate>MAX_RATE_INC*avg_rate)
                                                 dbg("HookReceiverP","Too many messages\n");

                                if(GATHER_TIME>DETECT_TIME)
					avg_rate = (GATHER_TIME/DETECT_TIME)*tmp_rate;
                                	else
					avg_rate = tmp_rate;

				tmp_rate = 0;

                        }

			for(i =0 ;i<sizeTmpList;i++)
				{
				for (j=0;j<16;j++)
					if( (((tmpList[i]>>j)&1)==1)&&(((blackFlag[i]>>j)&1)==0))
						{
						neighbours[i] = neighbours[i]| (1<<j); 
						/* if at least a new neighbor is added a new average rate needs to be calculated */
						rateTime = TRUE;
						}			
				tmpList[i] = 0 ;
				}
			sizeTmpList = 0;

			 
			/******print DETECT info***********/
/*                       dbg("HookReceiverP","Detect_Neighbours:\n");

			 for(i =0 ;i<sizeNeighbours;i++)
                                for (j=0 ;j<16;j++)
                                        if(neighbours[i]&(1<<j))
                                                dbg("HookReceiverP"," %d\n",i*16 + j);

                         dbg("HookReceiverP","Average size: %d \n",avg_size);
                         dbg("HookReceiverP","Average rate: %d \n",avg_rate);
*/
			call GDTimer.startOneShot(DETECT_TIME);
		break;
                }
        }
		
	 /* checks to se if a message fits the network characteristics */
        int IDSPass(message_t *msg , hook_msg_t *hm){
        
	        if(((double)hm->size >ALFA_SIZE*avg_size)||((double)hm->size<BETA_SIZE * avg_size))
                   {

			dbg("HookReceiverP","[INTRUSION]Received packet with out of range size  = %d\n",hm->size);

			return 0;
			}
                if(messageNr>ALFA_RATE*avg_rate)
                        {

			dbg("HookReceiverP","[INTRUSION]Exceeded the current transmision rate\n");
			return 0;
			}
                return 1;
        }

	 /* adds a node either to the neighbours list or the black flagged list */

        void addNodeToList(am_addr_t node,uint8_t type){

		/*
		the node lists are created to maintain info for about 16 different motes 
		in one element threw the its bit setting
		*/
                switch(type){
                case 0:
			blackFlag[node/16] = blackFlag[node/16]|(1 << (node%16)); 
               		dbg("HookReceiverP","A node has been black flaggedd: %d \n",node);

		 break;
                case 1:

			tmpList[node/16] = tmpList[node/16] | (1 << (node%16));
                        if(sizeTmpList<(node/16+1))	
				sizeTmpList = node/16 +1;
			dbg("HookReceiverP","A new tmp neighbour hass been added: %d \n",node);
                break;
                }

        }

	event message_t* ReceiveM.receive(message_t* msg, void* payload, uint8_t len){		
        
		hook_msg_t *hm;
		int src;
		bool flag= FALSE;
        	/*if dealing with a packet of different format then the standard one it is droped*/

        	if(len!=sizeof(hook_msg_t)){
                	dbg("HookReceiverP","Received a different packet type");
                	return msg;
        	}


        	switch(state){
        	case NO_STATE :
                	/*
                	start Timer
                	change state of hook
                	if message is for me and the header is the network header than add packet to vector for future processing
                	*/
                	state = GATHER ;

                	hm = (hook_msg_t*)payload;

                	if((hm!=NULL)&&(call AMPacket.isForMe(msg))&&(hm->header == HEADER)){
                        	memcpy(&messageVec[messageNr],msg,sizeof(message_t)+1);
                        	sizeVec[messageNr] = hm->size;
                        	messageNr = messageNr%MAX_MESSAGES +1;
                	}

			call GDTimer.startOneShot(GATHER_TIME);

                	break; 
        	case GATHER :
                	/*
                	if packet is for me add and the header is the network header it is added to vector for future processing
                	*/
                	hm = (hook_msg_t*)payload;
                	if((hm != NULL)&&(call AMPacket.isForMe(msg))&&(hm->header == HEADER)){
                        	memcpy(&messageVec[messageNr],msg,sizeof(message_t)+1);
                        	sizeVec[messageNr] = hm->size;
                        	messageNr = messageNr%MAX_MESSAGES +1;
                	}

                	break;
        	case DETECT :
                	/*dealing with a new node */
                	hm = (hook_msg_t*)payload;
		

                	if(hm == NULL)
				{
				dbg("HookReceiverP","Received a null packet \n");
				break;
				}
			if(!call AMPacket.isForMe(msg))
				{
				dbg("HookReceiverP","[INTRUSION]Received packet not for me msg=%d\n",hm->msg[0]);
				break;
				}
			if(hm->header!=HEADER)
				{
				 dbg("HookReceiverP","[INTRUSION]Received packet with wrong HEADER msg = %d\n",hm->msg[0]);
				break;
				}
			 src = call AMPacket.source(msg);

			if(((src/16)+1)>MAX_BLACK_FLAGS)
				{
				 dbg("HookReceiverP","[INTRUSION]Received packet from source with out of range ID = %d\n",src);
				 break;
				}
 
			if(((blackFlag[src/16]>>(src%16))&1) == 1)
				{
				 dbg("HookReceiverP","[INTRUSION]Received packet from black flag = %d\n",src);
                                break;

				}			
			flag = FALSE ;
			if(((neighbours[src/16]>>(src%16))&1) == 1)
                                {
                                 flag = TRUE ;

                                }


			if((!flag)&&(hm->size < ALFA_SIZE*avg_size)&&(hm->size>BETA_SIZE * avg_size)){
                                        addNodeToList(src,1);
                                        break;
                                }
			else
                        
			{
			if(!flag)
				{
				addNodeToList(src,0);
				break;
				}
			}
			
			
			/* dealing with a regular packet*/
                       	if(IDSPass(msg,hm)){
                             	messageNr++;
                             	signal Receive.receive(msg, payload, len);
				if(rateTime)
					tmp_rate++;
				break;
                             	
                		}
                       break; 
        	}
        	return msg;
	}

	//Packet
	command void Packet.clear(message_t* msg){
		
		return call PacketM.clear(msg);
	}
	
	command uint8_t Packet.payloadLength(message_t* msg){
		return call PacketM.payloadLength(msg);
	}
	
	command void Packet.setPayloadLength(message_t* msg, uint8_t len){
		return call PacketM.setPayloadLength(msg,len);
	}
	
	command uint8_t Packet.maxPayloadLength(){
		return call Packet.maxPayloadLength();
	}
	
	command void* Packet.getPayload(message_t* msg, uint8_t len){
		return call PacketM.getPayload(msg,len);
	}

	//AMPacket
	command am_addr_t AMPacket.address(){
		return call AMPacketM.address();
	}
	
	command am_addr_t AMPacket.destination(message_t* amsg){
		return call AMPacketM.destination(amsg);
	}
	
	command am_addr_t AMPacket.source(message_t* amsg){
		return call AMPacketM.source(amsg);
	}
	
	command void AMPacket.setDestination(message_t* amsg, am_addr_t addr){
		return call AMPacketM.setDestination(amsg, addr);
	}
	
	command void AMPacket.setSource(message_t* amsg, am_addr_t addr){
		return call AMPacketM.setSource(amsg, addr);
	}
	
	command bool AMPacket.isForMe(message_t* amsg){
		return call AMPacketM.isForMe(amsg);
	}
	
	command am_id_t AMPacket.type(message_t* amsg){
		return call AMPacketM.type(amsg);
	}
	
	command void AMPacket.setType(message_t* amsg, am_id_t t){
		return call AMPacketM.setType(amsg, t);
	}
	
	command am_group_t AMPacket.group(message_t* amsg){
		return call AMPacketM.group(amsg);
	}
	
	command void AMPacket.setGroup(message_t* amsg, am_group_t grp){
		return call AMPacketM.setGroup(amsg, grp);
	}
	
	command am_group_t AMPacket.localGroup(){
		return call AMPacketM.localGroup();
	}
		
}
